package com.example.mediacodecpro;

import android.graphics.ImageFormat;
import android.media.MediaCodecInfo;

import static android.graphics.ImageFormat.NV21;

public enum YUVConverter {
    NV21_YUV420p(ImageFormat.NV21, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar),
    NV21_YUV420sp(ImageFormat.NV21, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar),
    YV12_YUV420p(ImageFormat.YV12, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar),
    YV12_YUV420sp(ImageFormat.YV12, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar),
    ;

    private int cameraPreviewFormat;
    private int mediaCodecColorFormat;
    private boolean rotate90 = false;
    private boolean useLibyuv = false;
    private byte[] buf;

    YUVConverter(int cameraPreviewFormat, int mediaCodecColorFormat) {
        this.cameraPreviewFormat = cameraPreviewFormat;
        this.mediaCodecColorFormat = mediaCodecColorFormat;
    }

    public int getCameraPreviewFormat() {
        return cameraPreviewFormat;
    }

    public int getMediaCodecColorFormat() {
        return mediaCodecColorFormat;
    }

    public boolean isRotate90() {
        return rotate90;
    }

    public void setRotate90(boolean rotate90) {
        this.rotate90 = rotate90;
    }

    public boolean isUseLibyuv() {
        return useLibyuv;
    }

    public void setUseLibyuv(boolean useLibyuv) {
        this.useLibyuv = useLibyuv;
    }

    public void convert(byte[] input, byte[] output, int width, int height) {
        if (buf == null || buf.length != output.length) {
            buf = new byte[output.length];
        }
        switch (cameraPreviewFormat) {
            case ImageFormat.NV21:
                switch (mediaCodecColorFormat) {
                    case MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar:
                        if (useLibyuv) {
                            if (rotate90) {
                                YuvUtils.NV21ToI420andRotate90Clockwise(input, output, width, height);
                            } else {
                                YuvUtils.NV21ToI420(input, output, width, height);
                            }
                        } else {
                            if (rotate90) {
                                swapNV21toI420(input, buf, width, height);
                                YuvUtils.RotateI420(buf, output, width, height, 90);
                            } else {
                                swapNV21toI420(input, output, width, height);
                            }
                        }
                        break;
                    case MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar:
                        if (rotate90) {
                            swapNV21toNV12(input, buf, width, height);
                            YUV420spRotate90Clockwise(buf, output, width, height);
                        } else {
                            swapNV21toNV12(input, output, width, height);
                        }
                        break;
                    default:
                        throw new RuntimeException("unsupport color format for NV21");
                }
                break;
            case ImageFormat.YV12:
                switch (mediaCodecColorFormat) {
                    case MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar:
                        if (rotate90) {
                            swapYV12toI420(input, buf, width, height);
                            YuvUtils.RotateI420(buf, output, width, height, 90);
                        } else {
                            swapYV12toI420(input, output, width, height);
                        }
                        break;
                    case MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar:
                        if (rotate90) {
                            swapYV12toNV12(input, buf, width, height);
                            YUV420spRotate90Clockwise(buf, output, width, height);
                        } else {
                            swapYV12toNV12(input, output, width, height);
                        }
                        break;
                    default:
                        throw new RuntimeException("unsupport color format for YV12");
                }
                break;
        }
    }

    //    I420: YYYYYYYY UUVV ->YUV420P
    //    YV12: YYYYYYYY VVUU ->YUV420P
    //    NV12: YYYYYYYY UVUV ->YUV420SP
    //    NV21: YYYYYYYY VUVU ->YUV420SP

    //VVUU => UVUV
    private void swapYV12toNV12(byte[] yv12bytes, byte[] nv12bytes, int width, int height) {
        System.arraycopy(yv12bytes, 0, nv12bytes, 0, width * height);
        int size = width * height;

        for (int i = 0; i < size / 4; i++) {
            nv12bytes[size + 2 * i + 1] = yv12bytes[size + i];
            nv12bytes[size + 2 * i] = yv12bytes[size + size / 4 + i];
        }
    }

    //VVUU => UUVV
    private void swapYV12toI420(byte[] yv12bytes, byte[] i420bytes, int width, int height) {
        System.arraycopy(yv12bytes, 0, i420bytes, 0, width * height);
        System.arraycopy(yv12bytes, width * height + width * height / 4, i420bytes, width * height, width * height / 4);
        System.arraycopy(yv12bytes, width * height, i420bytes, width * height + width * height / 4, width * height / 4);
    }

    //VUVU => UVUV
    private void swapNV21toNV12(byte[] nv21bytes, byte[] i420bytes, int width, int height) {
        System.arraycopy(nv21bytes, 0, i420bytes, 0, width * height);
        int size = width * height;
        for (int i = size; i < size + size / 2; i += 2) {
            i420bytes[i] = nv21bytes[i + 1];
            i420bytes[i + 1] = nv21bytes[i];
        }
    }

    //VUVU => UUVV
    private void swapNV21toI420(byte[] nv21bytes, byte[] i420bytes, int width, int height) {
        System.arraycopy(nv21bytes, 0, i420bytes, 0, width * height);
        int size = width * height;
        for (int i = 0; i < size / 4; i += 1) {
            i420bytes[size + i] = nv21bytes[size + 2 * i + 1];
            i420bytes[size + size / 4 + i] = nv21bytes[size + 2 * i];
        }
    }

    //顺时针旋转90
    private void YUV420spRotate90Clockwise(byte[] src, byte[] dst, int width, int height) {
        int wh = width * height;
        int k = 0;
        for (int i = 0; i < width; i++) {
            for (int j = height - 1; j >= 0; j--) {
                dst[k] = src[width * j + i];
                k++;
            }
        }
        for (int i = 0; i < width; i += 2) {
            for (int j = height / 2 - 1; j >= 0; j--) {
                dst[k] = src[wh + width * j + i];
                dst[k + 1] = src[wh + width * j + i + 1];
                k += 2;
            }
        }
    }

    public int getPreviewBufferSize(int width, int height) {
        int size = 0;
        switch (cameraPreviewFormat) {
            case ImageFormat.YV12: {
                int yStride = (int) Math.ceil(width / 16.0) * 16;
                int uvStride = (int) Math.ceil((yStride / 2) / 16.0) * 16;
                int ySize = yStride * height;
                int uvSize = uvStride * height / 2;
                size = ySize + uvSize * 2;
            }
            break;

            case ImageFormat.NV21: {
                float bytesPerPix = (float) ImageFormat.getBitsPerPixel(cameraPreviewFormat) / 8;
                size = (int) (width * height * bytesPerPix);
            }
            break;
        }

        return size;
    }
}
